import cv2
import os
import numpy as np
from PIL import Image
from pdf2image import convert_from_path
from skimage.metrics import structural_similarity
from tqdm import tqdm

# 路径配置
basedir = os.path.abspath(os.path.dirname(__file__))  #basedir 变量就保存了当前脚本文件所在目录的绝对路径。
# 设置视频文件路径
video_path = "D:/2023DCProject/video/source/040910s.mp4"

# 设置截图保存路径和去重后截图保存路径
image_folder = f"{os.path.splitext(video_path)[0]}_frames"
unique_folder = f"{os.path.splitext(video_path)[0]}_unique_frames"

# 创建保存截图和去重后截图的文件夹
if not os.path.exists(image_folder):
    os.makedirs(image_folder)
if not os.path.exists(unique_folder):
    os.makedirs(unique_folder)

# 逐帧读取视频并保存截图
cap = cv2.VideoCapture(video_path)
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
for i in tqdm(range(frame_count)):
    ret, frame = cap.read()
    if ret:
        cv2.imwrite(f"{image_folder}/frame_{i:05d}.jpg", frame)
    else:
        break

# 遍历所有截图，计算相似度并去重
unique_images = []
for filename in tqdm(os.listdir(image_folder)):
    img = cv2.imread(os.path.join(image_folder, filename))
    if img is not None:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        if len(unique_images) == 0:
            unique_images.append((filename, img))
        else:
            found_duplicate = False
            for _, unique_img in unique_images:
                score, _ = structural_similarity(unique_img, img, full=True)
                if score > 0.9:
                    found_duplicate = True
                    break
            if not found_duplicate:
                unique_images.append((filename, img))

# 保存去重后的截图
for filename, img in tqdm(unique_images):
    cv2.imwrite(f"{unique_folder}/{filename}", img)

# 将去重后的截图按顺序拼接成PDF文档
images = []
for filename in sorted(os.listdir(unique_folder)):
    img = Image.open(os.path.join(unique_folder, filename))
    img = img.convert("RGB")
    images.append(img)

pdf_path = f"{os.path.splitext(video_path)[0]}.pdf"
images[0].save(pdf_path, save_all=True, append_images=images[1:])
